import random
from traceback import format_exc
import numpy
from ner.ner_case import NerCase
from ner.ner_model import NerModel
from static_word_vec.char_word_vec_encoder import CharWordVecEncoder
from static_word_vec.glyph_word_vec_encoder import GlyphWordVecEncoder
from static_word_vec.stroke_word_vec_encoder import StrokeWordVecEncoder
from static_word_vec.pinyin_word_vec_encoder import PinyinWordVecEncoder
from static_word_vec.word_vec_encoder_set import WordVecEncoderSet
from ner.dataset import Dataset
from ner.trainer import Trainer
import torch

if __name__ == '__main__':    

    # 固定随机数种子, 便于复现
    random.seed(1)
    numpy.random.seed(1)
    torch.manual_seed(1) 

    # 语料位置与类型，仅用于统计标签所用
    train_path = './data/train/'
    corpus_type = 'ccks'

    # 预训练模型的名字
    pretrained_model_name = 'hfl/chinese-roberta-wwm-ext-large' 

    # 待使用的已训练模型
    resume_from = './saves/-6734881434110966973_hfl-chinese-roberta-wwm-ext-large-dynamicfusion1x-bilstm1x-crf_epoch39_2022-09-10_08-20-53_44242.pt'

    # 要选用的其他特征词向量
    word_vector_encoder_manager = WordVecEncoderSet(
        #CharWordVecEncoder(),
        #PinyinWordVecEncoder(),
        #StrokeWordVecEncoder(),
        #ImageWordVecEncoder()
    )

    # 创建一个 NER 模型案例
    case = NerCase(pretrained_model_name = pretrained_model_name, # 使用预训练语言模型
                    use_dynamic_fusion = True, # 动态融合
                    num_dynamic_fusion_layers = 1,
                    word_vector_encoder_set = word_vector_encoder_manager, # 使用预先训练的特征
                    use_bilstm = True, # 使用 BiLSTM
                    num_bilstm_layers = 1, # BiLSTM 层数
                    use_crf = True, # 使用 CRF
                    resume_from = resume_from # 从已训练模型载入
                    )

    # 一个批次的待处理文本
    # 需要注意，文字长度不能超过510字，不然会超出处理能力
    # 批次不能过大，不然会超出显存限制
    batch_text = [
        '患者于2010年10月因\\\"上腹痛伴大便习惯改变\\\"外院行肠镜，活检病理示中分化腺癌；，外院B超示：'\
            '盆腔内2个巨大团块，考虑卵巢来源可能性大；盆腔大量积液。来我院就诊，考虑乙状结肠癌伴不完全梗阻、盆腔肿物，'\
            '2010-10-25我院行\\\"子宫切除+两附件切除+DIXON术+大网膜切除术开腹恶性肿瘤特殊治疗术\\\"，'\
            '术中5-FU1000MG肠腔化疗，门静脉5-FU250MG化疗，中人氟安800MG腹腔化疗，，病理回报：'\
            '乙状结肠中分化腺癌，两侧卵巢转移瘤，中央组淋巴结转移4/15，PT4BN2M1。'\
            '2010-11-10我院复查CT提示肝S5、6段肝转移。遂于2010-11-16至2011-3-11行FORFIRI方案化疗7程，'\
            '4程、6程复查CT疗效评价为SD，后于2011-4-6转外科行\\\"肝转移瘤切除术\\\"，，'\
            '术后病理：符合肠癌肝转移。后于2011-05-20至2011-08-05继续FORFIRI方案化疗4程。',

        '患者自1969年始，每遇天气转冷，即有咳嗽咯痰，清晨咳嗽较剧，痰量少，白色粘稠。' \
            '无气急、气喘、咯血及盗汗。每次持续约7～10天，经服四环素、枸橼酸维静宁（咳必清）等药即可好转。' \
            '每年发作2～3次，多在秋末冬初时。工作、生活不受影响。1981年以来，咳嗽、咯痰加重，早晚尤剧，有时伴气急。'\
            '每日痰量约10～20ml，为白色泡沫样。红青霉素等药物治疗可缓解。每次持续约3个月以上，天气转暖时上述症状缓解。',
    ]

    # 推断结果
    batch_result = case.predict_batch(batch_text,   # 批次文本
                        train_corpus_path = train_path,    # 仅用于统计标签，用训练集得到的标签与模型相同，其他数据集可能有所不同
                        corpus_type = corpus_type,  # 仅用于统计标签
                        return_style = 'parallel'   # 输出“原文与标签一一对应”的输出风格
                        )
    # 输出结果
    for result in batch_result:
        print(result)